import torch
import numpy as np
import matplotlib.pyplot as plt

from Config import Config
from train.metrics import MCR_MSE
from nets.load_net import gnn_model
from data.data_set import mRNADataset, gpu_setup

"""
读取训练得到的模型权重
并查看模型训练的效果
注意读取时要对Config文件进行修改
"""


def eval_network(model, device, graph, label, config):
    # 将某个数据放入模型中测试
    model.eval()
    with torch.no_grad():
        # 获取图的信息
        mask = graph.ndata['mask'].to(device)
        x = graph.ndata['nfeat'].to(device)
        e = graph.edata['efeat'].to(device)
        label = label.to(device)

        if config.add_loop_type == True:
            loop = graph.ndata['loop_type'].to(device)
            if config.pos_enc == True:  # 若使用循环类型则加入输入
                pos_enc = graph.ndata['pos_enc'].to(device)
                # sign_flip = torch.rand(batch_pos_enc.size(1)).to(device)
                # sign_flip[sign_flip >= 0.5] = 1.0
                # sign_flip[sign_flip < 0.5] = -1.0
                # batch_pos_enc = batch_pos_enc * sign_flip.unsqueeze(0)
                score = model.forward(graph, x, e, loop_type=loop, h_pos_enc=pos_enc)
            else:
                score = model.forward(graph, x, e, loop_type=loop)
        else:
            if config.pos_enc == True:  # 若使用循环类型则加入输入
                pos_enc = graph.ndata['pos_enc'].to(device)
                # sign_flip = torch.rand(batch_pos_enc.size(1)).to(device)
                # sign_flip[sign_flip >= 0.5] = 1.0
                # sign_flip[sign_flip < 0.5] = -1.0
                # batch_pos_enc = batch_pos_enc * sign_flip.unsqueeze(0)
                score = model.forward(graph, x, e, h_pos_enc=pos_enc)
            else:
                score = model.forward(graph, x, e)
        # loss = model.loss(score, label)
        # test_loss = loss.detach().item()
        test_assess = MCR_MSE(score, label, mask)
        print("MCRMSE: ", test_assess.cpu().numpy())

    label = label.cpu().numpy()
    score = score.cpu().numpy()
    return label, score


def plat_res(label, score, seq_scored):
    # 画图查看结果(真实标签, 预测分数, 需要预测的长度)
    length, col = label.shape
    x = range(seq_scored)
    y_labels = ['reactivity', 'deg_Mg_pH10', 'deg_Mg_50C']
    for i in range(col):
        plt.figure(i)
        y1 = label[:seq_scored, i]
        y2 = score[:seq_scored, i]
        plt.plot(x, y1, label='label', color='blue')
        plt.plot(x, y2, label='score', color='red')
        plt.xlabel('position')
        plt.ylabel(y_labels[i])
        plt.legend(loc='best')
    plt.show()


if __name__ == '__main__':
    model_name = 'GatedGCN'
    data_file = "../dataset/valid.json"
    seq_scored = 68  # 按要求只看前68个碱基的效果
    """GPU设置"""
    device = gpu_setup(True, 0)
    config = Config(model_name, device)
    data_set = mRNADataset(data_file, config, mode='test')

    id = 10  # 查看哪条数据
    graph = data_set.graph_list[id]
    label = data_set.node_labels[id]
    label = torch.tensor(label)
    # 设置位置编码
    # data_set._add_positional_encodings(pos_enc_dim=config.pos_enc_dim)

    model = gnn_model(model_name, config)
    model = model.to(device)
    model.load_state_dict(torch.load('../checkpoints/GatedGCN/best.pkl', map_location=torch.device('cpu')))
    label, score = eval_network(model, device, graph, label, config)

    # print(label, score)
    plat_res(label, score, seq_scored)
